#include <bsl_err_internal.h>
#include <bsl_errno.h>
#include <bsl_params.h>
#include <crypt_algid.h>
#include <crypt_dh.h>
#include <crypt_errno.h>
#include <crypt_utils.h>
#include <securec.h>

#include <wd.h>
#include <wd_bmm.h>
#include <wd_dh.h>

#include "kaev1_dh.h"
#include "kaev1_dh_local.h"
#include "kaep_err.h"
#include "kaev1_memory.h"

#define DH_BLOCK_NUM  16
#define DH_BLOCK_SIZE 4096
#define CRYPT_DH_TRY_CNT_MAX 100

KaeQueuePoolHead *gDHNodePool = NULL;

int WdDHInitQnodePool(void)
{
    KaeQueuePoolDestroy(gDHNodePool);

    gDHNodePool = KaeInitQueuePool(WCRYPTO_DH);
    if (gDHNodePool == NULL) {
        return KAE_WD_DO_DIGEST_FAIL;
    }
    return CRYPT_SUCCESS;
}

int WdDHDestroyQnodePool(void)
{
    KaeQueuePoolDestroy(gDHNodePool);
    return CRYPT_SUCCESS;
}

static void HpreDhCallback(const void *message, void *tag)
{
    if (!message || !tag) {
        return;
    }
    struct wcrypto_dh_msg *msg = (struct wcrypto_dh_msg *)message;
    KAE_DH_Ctx *ctx = (KAE_DH_Ctx *)tag;
    ctx->opData.pri = msg->out;
    ctx->opData.pri_bytes = msg->out_bytes;
    ctx->opData.status = msg->result;
}

void KAE_DH_Free(KAE_DH_Ctx *ctx)
{
    if (ctx == NULL) {
        return;
    }

    KAE_DH_FreePara(ctx->para);

    if (ctx->dhCtx) {
        wcrypto_del_dh_ctx(ctx->dhCtx);
        ctx->dhCtx = NULL;
    }

    KaePutNodeToPool(gDHNodePool, ctx->entry);

    BN_Destroy(ctx->prvKey);
    BN_Destroy(ctx->pubKey);

    free(ctx);
}

KAE_DH_Ctx *KAE_DH_NewCtx()
{
    KAE_DH_Ctx *ctx = (KAE_DH_Ctx *)malloc(sizeof(KAE_DH_Ctx));
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }

    memset(ctx, 0, sizeof(KAE_DH_Ctx));

    ctx->entry = KaeGetNodeFromPool(gDHNodePool);
    if (ctx->entry == NULL) {
        KAE_DH_Free(ctx);
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }

    ctx->queue = ctx->entry->queue;
    ctx->pool = ctx->entry->mempool;
    return ctx;
}

int32_t ParaQCheck(BN_BigNum *q, BN_BigNum *r)
{
    // 1. Determine the length.
    if (BN_Bits(q) < DH_MIN_QBITS) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    // 2. Parity and even judgment
    if (BN_GetBit(q, 0) == 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    // 3. Compare q and r.
    if (BN_Cmp(q, r) >= 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    // 4. Check the pq multiple relationship.
    BN_Optimizer *opt = BN_OptimizerCreate();
    if (opt == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }
    int32_t ret = BN_Div(NULL, r, r, q, opt);
    BN_OptimizerDestroy(opt);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    // (p - 1) % q == 0
    if (!BN_IsZero(r)) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    return CRYPT_SUCCESS;
}

static int32_t ParaDataCheck(const DH_Para *para)
{
    int32_t ret;
    const BN_BigNum *p = para->p;
    const BN_BigNum *g = para->g;
    // 1. Determine the length.
    uint32_t pBits = BN_Bits(p);
    if (pBits < DH_MIN_PBITS || pBits > DH_MAX_PBITS) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    // 2. P parity and g value judgment
    // p is an odd number
    if (BN_GetBit(p, 0) == 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    // g != 0 && g != 1
    if (BN_IsZero(g) || BN_IsOne(g)) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    BN_BigNum *r = BN_Create(pBits + 1);
    if (r == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    // r = p - 1
    ret = BN_SubLimb(r, p, 1);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    // g < p - 1
    if (BN_Cmp(g, r) >= 0) {
        ret = CRYPT_DH_PARA_ERROR;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    if (para->q != NULL) {
        ret = ParaQCheck(para->q, r);
    }
ERR:
    BN_Destroy(r);
    return ret;
}

uint32_t DhGetBits(DH_Para *para)
{
    if (para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return 0;
    }

    uint32_t bits = BN_Bits(para->p);
    for (int32_t i = 0; i < DH_BITS_SIZE; i--) {
        if (bits < DH_BITS[i]) {
            return DH_BITS[i];
        }
    }

    return 0;
}

bool KAE_DH_IsG2(const BN_BigNum *a)
{
    return (BN_Bits(a) == 2 && !BN_IsOdd(a));
}

int32_t KAE_DH_SetParaEx(KAE_DH_Ctx *ctx, DH_Para *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    int32_t ret = ParaDataCheck(para);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    BN_Destroy(ctx->prvKey);
    BN_Destroy(ctx->pubKey);
    KAE_DH_FreePara(ctx->para);
    uint32_t bits = BN_Bits(para->p);
    ctx->prvKey = NULL;
    ctx->pubKey = NULL;
    ctx->para = para;

    ctx->setup.key_bits = bits;
    ctx->setup.is_g2 = KAE_DH_IsG2(ctx->para->g);
    ctx->setup.cb = (wcrypto_cb)HpreDhCallback;
    ctx->setup.is_g2 = 0;
    ctx->setup.br.alloc = (void *)wd_alloc_blk;
    ctx->setup.br.free = (void *)wd_free_blk;
    ctx->setup.br.iova_map = (void *)wd_blk_iova_map;
    ctx->setup.br.iova_unmap = (void *)wd_blk_iova_unmap;
    ctx->setup.br.usr = ctx->pool;

    return CRYPT_SUCCESS;
}

int32_t KAE_DH_SetPara(KAE_DH_Ctx *ctx, const BSL_Param *para)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    DH_Para *dhPara = KAE_DH_NewPara(para);
    if (dhPara == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_ERR_NEW_PARA_FAIL);
        return CRYPT_EAL_ERR_NEW_PARA_FAIL;
    }
    int32_t ret = ParaDataCheck(dhPara);
    if (ret != CRYPT_SUCCESS) {
        KAE_DH_FreePara(dhPara);
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }

    uint32_t bits = BN_Bits(dhPara->p);
    if (bits == 0) {
        KAE_DH_FreePara(dhPara);
        BSL_ERR_PUSH_ERROR(CRYPT_DH_CREATE_PARA_FAIL);
        return CRYPT_DH_CREATE_PARA_FAIL;
    }

    BN_Destroy(ctx->prvKey);
    BN_Destroy(ctx->pubKey);
    KAE_DH_FreePara(ctx->para);
    ctx->prvKey = NULL;
    ctx->pubKey = NULL;
    ctx->para = dhPara;

    ctx->setup.key_bits = bits;
    ctx->setup.is_g2 = KAE_DH_IsG2(ctx->para->g);
    ctx->setup.cb = (wcrypto_cb)HpreDhCallback;
    ctx->setup.is_g2 = 0;
    ctx->setup.br.alloc = (void *)wd_alloc_blk;
    ctx->setup.br.free = (void *)wd_free_blk;
    ctx->setup.br.iova_map = (void *)wd_blk_iova_map;
    ctx->setup.br.iova_unmap = (void *)wd_blk_iova_unmap;
    ctx->setup.br.usr = ctx->pool;

    return ret;
}

int32_t KAE_DH_SetParamById(KAE_DH_Ctx *ctx, CRYPT_PKEY_ParaId id)
{
    DH_Para *para = KAE_DH_NewParaById(id);
    if (para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_ERR_NEW_PARA_FAIL);
        return CRYPT_EAL_ERR_NEW_PARA_FAIL;
    }
    int32_t ret = KAE_DH_SetParaEx(ctx, para);
    if (ret != CRYPT_SUCCESS) {
        KAE_DH_FreePara(para);
    }
    return ret;
}

int32_t GetDhParamP(const DH_Para *para, BSL_Param *param)
{
    BSL_Param *temp = BSL_PARAM_FindParam(param, CRYPT_PARAM_DH_P);
    if (temp == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    temp->useLen = temp->valueLen;
    int32_t ret = BN_Bn2Bin(para->p, temp->value, &temp->useLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }

    return CRYPT_SUCCESS;
}

int32_t GetDhParamQ(const DH_Para *para, BSL_Param *param)
{
    BSL_Param *temp = BSL_PARAM_FindParam(param, CRYPT_PARAM_DH_Q);
    if (temp == NULL) {
        return CRYPT_SUCCESS;
    }

    if (para->q == NULL) {
        temp->useLen = 0;
        return CRYPT_SUCCESS;
    }

    temp->useLen = temp->valueLen;
    int32_t ret = BN_Bn2Bin(para->q, temp->value, &temp->useLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }

    return CRYPT_SUCCESS;
}

static int32_t GetDhParamG(const DH_Para *para, BSL_Param *param)
{
    BSL_Param *temp = BSL_PARAM_FindParam(param, CRYPT_PARAM_DH_G);
    if (temp == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    temp->useLen = temp->valueLen;
    int32_t ret = BN_Bn2Bin(para->g, temp->value, &temp->useLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }

    return CRYPT_SUCCESS;
}

int32_t KAE_DH_GetPara(const KAE_DH_Ctx *ctx, BSL_Param *param)
{
    if (ctx == NULL || param == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    int32_t ret = GetDhParamP(ctx->para, param);
    if (ret != CRYPT_SUCCESS) {
        return ret;
    }

    ret = GetDhParamQ(ctx->para, param);
    if (ret != CRYPT_SUCCESS) {
        return ret;
    }

    ret = GetDhParamG(ctx->para, param);
    if (ret != CRYPT_SUCCESS) {
        return ret;
    }

    return CRYPT_SUCCESS;
}

int32_t PrvLenCheck(const KAE_DH_Ctx *ctx, const BSL_Param *prv)
{
    if (ctx->para->q != NULL) {
        if (BN_Bytes(ctx->para->q) < prv->valueLen) {
            BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
            return CRYPT_DH_KEYINFO_ERROR;
        }
    } else {
        if (BN_Bytes(ctx->para->p) < prv->valueLen) {
            BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
            return CRYPT_DH_KEYINFO_ERROR;
        }
    }
    return CRYPT_SUCCESS;
}

int32_t GetXLimb(BN_BigNum *xLimb, const BN_BigNum *p, const BN_BigNum *q)
{
    if (q != NULL) {
        // xLimb = q - 1
        return BN_SubLimb(xLimb, q, 1);
    }
    // xLimb = p - 2
    return BN_SubLimb(xLimb, p, 2);
}

int32_t KAE_DH_SetPrvKey(KAE_DH_Ctx *ctx, const BSL_Param *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    const BSL_Param *prv = BSL_PARAM_FindConstParam(para, CRYPT_PARAM_DH_PRVKEY);
    if (prv == NULL || prv->value == NULL || prv->valueLen == 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    int32_t ret = PrvLenCheck(ctx, prv);
    if (ret != CRYPT_SUCCESS) {
        return ret;
    }
    BN_BigNum *bnX = BN_Create(BN_BYTES_TO_BITS(prv->valueLen));
    BN_BigNum *xLimb = BN_Create(BN_Bits(ctx->para->p) + 1);
    if (bnX == NULL || xLimb == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    ret = GetXLimb(xLimb, ctx->para->p, ctx->para->q);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    ret = BN_Bin2Bn(bnX, prv->value, prv->valueLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    // Satisfy x <= q - 1 or x <= p - 2
    if (BN_Cmp(bnX, xLimb) > 0) {
        ret = CRYPT_DH_KEYINFO_ERROR;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    // x != 0
    if (BN_IsZero(bnX)) {
        ret = CRYPT_DH_KEYINFO_ERROR;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    BN_Destroy(xLimb);
    BN_Destroy(ctx->prvKey);
    ctx->prvKey = bnX;
    return ret;
ERR:
    BN_Destroy(bnX);
    BN_Destroy(xLimb);
    return ret;
}

// No parameter information is required for setting the public key.
// Therefore, the validity of the public key is not checked during the setting.
// The validity of the public key is checked during the calculation of the
// shared key.
int32_t KAE_DH_SetPubKey(KAE_DH_Ctx *ctx, const BSL_Param *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    const BSL_Param *pub = BSL_PARAM_FindConstParam(para, CRYPT_PARAM_DH_PUBKEY);
    if (pub == NULL) {
        pub = BSL_PARAM_FindConstParam(para, CRYPT_PARAM_PKEY_ENCODE_PUBKEY);
    }

    if (pub == NULL || pub->value == NULL || pub->valueLen == 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (pub->valueLen > BN_BITS_TO_BYTES(DH_MAX_PBITS)) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }
    BN_BigNum *bnY = BN_Create(BN_BYTES_TO_BITS(pub->valueLen));
    if (bnY == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }
    int32_t ret = BN_Bin2Bn(bnY, pub->value, pub->valueLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    BN_Destroy(ctx->pubKey);
    ctx->pubKey = bnY;
    return ret;
ERR:
    BN_Destroy(bnY);
    return ret;
}

int32_t KAE_DH_GetPrvKey(const KAE_DH_Ctx *ctx, BSL_Param *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    BSL_Param *prv = BSL_PARAM_FindParam(para, CRYPT_PARAM_DH_PRVKEY);
    if (prv == NULL || prv->value == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (ctx->prvKey == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }
    if (ctx->para->q != NULL) {
        if (BN_Bytes(ctx->para->q) > prv->valueLen) {
            BSL_ERR_PUSH_ERROR(CRYPT_DH_BUFF_LEN_NOT_ENOUGH);
            return CRYPT_DH_BUFF_LEN_NOT_ENOUGH;
        }
    } else {
        if (BN_Bytes(ctx->para->p) > prv->valueLen) {
            BSL_ERR_PUSH_ERROR(CRYPT_DH_BUFF_LEN_NOT_ENOUGH);
            return CRYPT_DH_BUFF_LEN_NOT_ENOUGH;
        }
    }
    uint32_t useLen = prv->valueLen;
    int32_t ret = BN_Bn2Bin(ctx->prvKey, prv->value, &(useLen));
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    prv->useLen = useLen;
    return CRYPT_SUCCESS;
}

int32_t KAE_DH_GetPubKey(const KAE_DH_Ctx *ctx, BSL_Param *para)
{
    if (ctx == NULL || para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    BSL_Param *pub = BSL_PARAM_FindParam(para, CRYPT_PARAM_DH_PUBKEY);
    if (pub == NULL) {
        pub = BSL_PARAM_FindParam(para, CRYPT_PARAM_PKEY_ENCODE_PUBKEY);
    }

    if (pub == NULL || pub->value == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (ctx->pubKey == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }
    if (BN_Bytes(ctx->pubKey) > pub->valueLen) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_BUFF_LEN_NOT_ENOUGH);
        return CRYPT_DH_BUFF_LEN_NOT_ENOUGH;
    }
    uint32_t useLen = pub->valueLen;
    int32_t ret = BN_Bn2Bin(ctx->pubKey, pub->value, &(useLen));
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    pub->useLen = useLen;
    return CRYPT_SUCCESS;
}

static void RefreshKaeCtx(KAE_DH_Ctx *dhCtx, BN_BigNum *x, BN_BigNum *y, int32_t ret)
{
    if (ret == CRYPT_SUCCESS) {
        BN_Destroy(dhCtx->prvKey);
        BN_Destroy(dhCtx->pubKey);
        dhCtx->prvKey = x;
        dhCtx->pubKey = y;
    } else {
        BN_Destroy(x);
        BN_Destroy(y);
    }
}

int32_t KAE_DH_Gen(KAE_DH_Ctx *ctx)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    if (DH_CheckBitsLength(ctx->setup.key_bits) != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_ERR_KEY_BITS);
        return KAE_WD_ERR_KEY_BITS;
    }

    ctx->dhCtx = wcrypto_create_dh_ctx(ctx->queue, &ctx->setup);
    if (ctx->dhCtx == NULL) {
        BSL_ERR_PUSH_ERROR(KAE_WD_CREATE_CTX_FAIL);
        return KAE_WD_CREATE_CTX_FAIL;
    }

    int keySize = BN_BITS_TO_BYTES(ctx->setup.key_bits);
    unsigned char *agBin = (unsigned char *)wd_alloc_blk(ctx->pool);
    if (agBin == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }
    unsigned char *prvBin = agBin + keySize;
    unsigned char *apBin = prvBin + keySize;
    memset(agBin, 0, keySize * 4);

    int32_t ret;
    uint32_t bits = BN_Bits(ctx->para->p);
    BN_BigNum *x = BN_Create(bits);
    BN_BigNum *y = BN_Create(bits);
    BN_BigNum *minP = BN_Create(bits);
    BN_BigNum *xLimb = BN_Create(bits);

    if (x == NULL || y == NULL || minP == NULL || xLimb == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }

    ret = BN_SubLimb(minP, ctx->para->p, 1);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    ret = GetXLimb(xLimb, ctx->para->p, ctx->para->q);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }

    struct wd_dtb gtb;
    gtb.dsize = keySize;

    ret = BN_Bn2Bin(ctx->para->g, agBin, &gtb.dsize);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }
    gtb.data = (char *)agBin;
    gtb.bsize = keySize;
    if (wcrypto_set_dh_g(ctx->dhCtx, &gtb) != 0) {
        ret = KAE_DH_SET_G_ERROR;
        BSL_ERR_PUSH_ERROR(ret);
        goto EXIT;
    }

    for (size_t cnt = 0; cnt < CRYPT_DH_TRY_CNT_MAX; cnt++) {
        /*  Generate private key x for [1, q-1] or [1, p-2] */
        ret = BN_RandRange(x, xLimb);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }
        ret = BN_AddLimb(x, x, 1);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }

        uint32_t pbytes = keySize;
        uint32_t xbytes = keySize;

        ret = BN_Bn2Bin(ctx->para->p, apBin, &pbytes);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }

        ret = BN_Bn2Bin(x, prvBin, &xbytes);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }

        ctx->opData.pbytes = pbytes;
        ctx->opData.xbytes = xbytes;

        ctx->opData.x_p = prvBin;
        ctx->opData.pri = apBin + keySize;
        ctx->opData.op_type = WCRYPTO_DH_PHASE1;
        if (wcrypto_do_dh(ctx->dhCtx, &ctx->opData, NULL) != 0) {
            ret = KAE_DH_GEN_ERROR;
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }
        const unsigned char *pubBin = (const unsigned char *)ctx->opData.pri;
        if (pubBin == NULL) {
            ret = KAE_DH_GEN_ERROR;
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }

        ret = BN_Bin2Bn(y, pubBin, ctx->opData.pri_bytes);
        if (ret != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(ret);
            goto EXIT;
        }

        if (BN_IsZero(y) || BN_IsOne(y) || BN_Cmp(y, minP) >= 0) {
            continue;
        }
        goto EXIT;
    }
    ret = CRYPT_DH_RAND_GENERATE_ERROR;
    BSL_ERR_PUSH_ERROR(ret);
EXIT:
    RefreshKaeCtx(ctx, x, y, ret);
    BN_Destroy(minP);
    BN_Destroy(xLimb);

    if (ctx->pool && agBin) {
        ctx->setup.br.free(ctx->pool, agBin);
        agBin = NULL;
    }
    if (ctx->dhCtx) {
        wcrypto_del_dh_ctx(ctx->dhCtx);
        ctx->dhCtx = NULL;
    }
    return ret;
}

static int32_t PubCheck(const BN_BigNum *y, const BN_BigNum *minP, const BN_BigNum *q, BN_Mont *montP,
                        BN_Optimizer *opt)
{
    // y != 0, y != 1
    if (BN_IsZero(y) || BN_IsOne(y)) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }
    // y < p - 1
    if (BN_Cmp(y, minP) >= 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }
    if (q == NULL) {
        return CRYPT_SUCCESS; // The parameter q does not exist, this function is ended early.
    }
    // Verify q.
    BN_BigNum *r = BN_Create(BN_Bits(minP));
    if (r == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }
    int32_t ret = BN_MontExp(r, y, q, montP, opt);
    if (ret != CRYPT_SUCCESS) {
        BN_Destroy(r);
        BSL_ERR_PUSH_ERROR(ret);
        return ret;
    }
    if (!BN_IsOne(r)) {
        BN_Destroy(r);
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }
    BN_Destroy(r);
    return CRYPT_SUCCESS;
}

static void CheckAndFillZero(uint8_t *shareKey, uint32_t *shareKeyLen, uint32_t bytes)
{
    int32_t i;
    if (*shareKeyLen == bytes) { // (*shareKeyLen > bytes) is not possible
        return;
    }
    uint32_t fill = bytes - *shareKeyLen;
    for (i = (int32_t)*shareKeyLen - 1; i >= 0; i--) {
        shareKey[i + (int32_t)fill] = shareKey[i];
    }
    for (i = 0; i < (int32_t)fill; i++) {
        shareKey[i] = 0;
    }
    *shareKeyLen = bytes;
}

int32_t KAE_DH_ComputeShareKey(const KAE_DH_Ctx *ctx, const KAE_DH_Ctx *pubKey, uint8_t *out, uint32_t *outLen)
{
    if (ctx == NULL || pubKey == NULL || out == NULL || outLen == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }
    if (ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }
    if (ctx->prvKey == NULL || pubKey->pubKey == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }

    if (BN_Bytes(ctx->para->p) > *outLen) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_BUFF_LEN_NOT_ENOUGH);
        return CRYPT_DH_BUFF_LEN_NOT_ENOUGH;
    }

    if (DH_CheckBitsLength(ctx->setup.key_bits) != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_ERR_KEY_BITS);
        return KAE_WD_ERR_KEY_BITS;
    }

    KAE_DH_Ctx *ncCtx = (KAE_DH_Ctx *)ctx;
    void *dhCtx = wcrypto_create_dh_ctx(ncCtx->queue, &ncCtx->setup);
    if (dhCtx == NULL) {
        BSL_ERR_PUSH_ERROR(KAE_WD_CREATE_CTX_FAIL);
        return KAE_WD_CREATE_CTX_FAIL;
    }

    struct wcrypto_dh_op_data opData;
    int keySize = ctx->setup.key_bits >> 3;
    unsigned char *agBin = (unsigned char *)wd_alloc_blk(ctx->pool);
    if (agBin == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }

    int32_t ret;

    BN_BigNum *minP = BN_Create(BN_Bits(ctx->para->p) + 1);
    BN_BigNum *r = BN_Create(BN_Bits(ctx->para->p));
    BN_Mont *mont = BN_MontCreate(ctx->para->p);
    BN_Optimizer *opt = BN_OptimizerCreate();
    if (minP == NULL || r == NULL || mont == NULL || opt == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    ret = BN_SubLimb(minP, ctx->para->p, 1);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    /* Check whether the public key meets the requirements. */
    ret = PubCheck(pubKey->pubKey, minP, ctx->para->q, mont, opt);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }

    unsigned char *prvBin = agBin + keySize;
    unsigned char *apBin = prvBin + keySize;
    memset(agBin, 0, keySize * 4);

    struct wd_dtb gtb;
    gtb.dsize = keySize;
    GOTO_ERR_IF(BN_Bn2Bin(ctx->para->g, agBin, &gtb.dsize), ret);
    gtb.data = (char *)agBin;
    gtb.bsize = keySize;
    if (wcrypto_set_dh_g(dhCtx, &gtb) != 0) {
        ret = KAE_DH_SET_G_ERROR;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    opData.op_type = WCRYPTO_DH_PHASE2;
    opData.pbytes = keySize;
    opData.xbytes = keySize;
    opData.pvbytes = keySize;

    uint32_t pbytes = keySize;
    uint32_t xbytes = keySize;
    uint32_t pvbytes = keySize;

    GOTO_ERR_IF(BN_Bn2Bin(ctx->para->p, apBin, &pbytes), ret);
    GOTO_ERR_IF(BN_Bn2Bin(ctx->prvKey, prvBin, &xbytes), ret);
    GOTO_ERR_IF(BN_Bn2Bin(pubKey->pubKey, agBin, &pvbytes), ret);

    opData.pbytes = pbytes;
    opData.xbytes = xbytes;
    opData.pvbytes = pvbytes;

    opData.x_p = prvBin;
    opData.pri = apBin + keySize;
    opData.pv = agBin;
    if (wcrypto_do_dh(dhCtx, &opData, NULL) != 0) {
        ret = KAE_DH_GEN_SHARED_ERROR;
        BSL_ERR_PUSH_ERROR(ret);
        goto ERR;
    }
    memcpy(out, opData.pri, opData.pri_bytes);
    *outLen = opData.pri_bytes;
    CheckAndFillZero(out, outLen, BN_Bytes(ctx->para->p));
ERR:
    BN_Destroy(minP);
    BN_Destroy(r);
    BN_MontDestroy(mont);
    BN_OptimizerDestroy(opt);

    if (ctx->pool && agBin) {
        ctx->setup.br.free(ctx->pool, agBin);
        agBin = NULL;
    }

    if (dhCtx) {
        wcrypto_del_dh_ctx(dhCtx);
        dhCtx = NULL;
    }
    return ret;
}

int32_t CRYPT_DH_GetLen(const KAE_DH_Ctx *ctx, GetLenFunc func, void *val, uint32_t len)
{
    if (val == NULL || len != sizeof(int32_t)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    *(int32_t *)val = func(ctx);
    return CRYPT_SUCCESS;
}

CRYPT_PKEY_ParaId KAE_DH_GetParaId(const KAE_DH_Ctx *ctx)
{
    if (ctx == NULL || ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_ERR_ALGID);
        return CRYPT_PKEY_PARAID_MAX;
    }
    return ctx->para->id;
}

int32_t KAE_DH_GetBits(const KAE_DH_Ctx *ctx)
{
    if (ctx == NULL || ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_ERR_ALGID);
        return CRYPT_PKEY_PARAID_MAX;
    }
    return BN_Bits(ctx->para->p);
}

int32_t KAE_DH_GetSecBits(const KAE_DH_Ctx *ctx)
{
    if (ctx == NULL || ctx->para == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_EAL_ERR_ALGID);
        return CRYPT_PKEY_PARAID_MAX;
    }
    return DhGetBits(ctx->para);
}

int32_t KAE_DH_Ctrl(KAE_DH_Ctx *ctx, int32_t opt, void *val, uint32_t len)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    switch (opt) {
        case CRYPT_CTRL_GET_PARAID:
            return CRYPT_DH_GetLen(ctx, (GetLenFunc)KAE_DH_GetParaId, val, len);
        case CRYPT_CTRL_GET_BITS:
            return CRYPT_DH_GetLen(ctx, (GetLenFunc)KAE_DH_GetBits, val, len);
        case CRYPT_CTRL_GET_SECBITS:
            return CRYPT_DH_GetLen(ctx, (GetLenFunc)KAE_DH_GetSecBits, val, len);
        case CRYPT_CTRL_SET_PARA_BY_ID:
            return KAE_DH_SetParamById(ctx, *(CRYPT_PKEY_ParaId *)val);
        default:
            break;
    }
    BSL_ERR_PUSH_ERROR(CRYPT_DH_UNSUPPORTED_CTRL_OPTION);
    return CRYPT_DH_UNSUPPORTED_CTRL_OPTION;
}

DH_Para *KAE_DH_NewParaFromBits(uint32_t bits)
{
    DH_Para *newPara = (DH_Para *)malloc(sizeof(DH_Para));

    newPara->p = BN_Create(bits);
    newPara->q = BN_Create(bits);
    newPara->g = BN_Create(bits);

    if (newPara == NULL || newPara->p == NULL || newPara->q == NULL || newPara->g == NULL) {
        KAE_DH_FreePara(newPara);
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }

    return newPara;
}

static DH_Para *ParaDup(const DH_Para *para)
{
    DH_Para *ret = (DH_Para *)malloc(sizeof(DH_Para));
    if (ret == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    ret->p = BN_Dup(para->p);
    ret->q = BN_Dup(para->q);
    ret->g = BN_Dup(para->g);
    ret->id = para->id;
    if (ret->p == NULL || ret->g == NULL) {
        KAE_DH_FreePara(ret);
        BSL_ERR_PUSH_ERROR(CRYPT_DH_CREATE_PARA_FAIL);
        return NULL;
    }
    if (para->q != NULL && ret->q == NULL) {
        KAE_DH_FreePara(ret);
        BSL_ERR_PUSH_ERROR(CRYPT_DH_CREATE_PARA_FAIL);
        return NULL;
    }
    return ret;
}

KAE_DH_Ctx *KAE_DH_DupCtx(KAE_DH_Ctx *ctx)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return NULL;
    }

    KAE_DH_Ctx *newCtx = KAE_DH_NewCtx();
    if (newCtx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }

    GOTO_ERR_IF_SRC_NOT_NULL(newCtx->prvKey, ctx->prvKey, BN_Dup(ctx->prvKey), CRYPT_MEM_ALLOC_FAIL);
    GOTO_ERR_IF_SRC_NOT_NULL(newCtx->pubKey, ctx->pubKey, BN_Dup(ctx->pubKey), CRYPT_MEM_ALLOC_FAIL);
    GOTO_ERR_IF_SRC_NOT_NULL(newCtx->para, ctx->para, ParaDup(ctx->para), CRYPT_MEM_ALLOC_FAIL);
    return newCtx;

ERR:
    KAE_DH_Free(newCtx);
    return NULL;
}

int32_t KAE_DH_Cmp(const KAE_DH_Ctx *a, const KAE_DH_Ctx *b)
{
    if (a == NULL || b == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (a->pubKey == NULL || b->pubKey == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_KEYINFO_ERROR;
    }

    if (BN_Cmp(a->pubKey, b->pubKey) != 0) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_KEYINFO_ERROR);
        return CRYPT_DH_PUBKEY_NOT_EQUAL;
    }

    if ((a->para == NULL) != (b->para == NULL)) {
        BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_ERROR);
        return CRYPT_DH_PARA_ERROR;
    }

    if (a->para != NULL) {
        if (BN_Cmp(a->para->p, b->para->p) != 0 || BN_Cmp(a->para->q, b->para->q) != 0 ||
            BN_Cmp(a->para->g, b->para->g) != 0) {
            BSL_ERR_PUSH_ERROR(CRYPT_DH_PARA_NOT_EQUAL);
            return CRYPT_DH_PARA_NOT_EQUAL;
        }
    }
    return CRYPT_SUCCESS;
}